import numpy as np
import matplotlib

matplotlib.use(backend="TkAgg")
import matplotlib.pyplot as plt
import pandas as pd
from math import isclose
import random
from collections import Counter
import numpy.random as npr
import numpy.linalg as linalg

np.set_printoptions(precision=4, suppress=True)

print("Section7: Reversible chain and random walk on undirected graph")
# construct adjacency for undirected graph and transition matrix for simple random walk

adj = np.array([
    [0, 1, 1, 0],
    [1, 0, 1, 1],
    [1, 1, 0, 1],
    [0, 1, 1, 0]
])
degrees = adj.sum(axis=1)
P_rw = (adj.T / degrees).T  # row-normalize

# 度数归一化得到转移概率
# P_rw[0] = [0, 1/2, 1/2, 0]    # 从节点0: 50%到1, 50%到2
# P_rw[1] = [1/3, 0, 1/3, 1/3]  # 从节点1: 均等概率到0,2,3
print(pd.DataFrame(P_rw))
print("degrees:", degrees)
# stationary distribution proportional to degrees
# 对于无向图上的简单随机游走，平稳分布与度数成正比：
# π_i = deg(i) / (2 × 边数)
# 这里总度数为10（每条边被计算两次），所以：
# s_rw = [2/10, 3/10, 3/10, 2/10] = [0.2, 0.3, 0.3, 0.2]
s_rw = degrees / degrees.sum()
print("stationary (theoretical) proportional to degrees:", s_rw)


# verify detailed balance s_i*q_ij = s_j*q_ji for all i,j
def check_detailed_balance(P, s, tol=1e-10):
    n = len(s)
    violations = []
    for i in range(n):
        for j in range(n):
            # 细致平衡条件:
            # π_i × P_ij = π_j × P_ji  (对所有i,j)
            # 节点0→1: 0.2 × 0.5 = 0.1
            # 节点1→0: 0.3 × 0.333... = 0.1
            # 两者相等，满足细致平衡
            if not isclose(s[i] * P[i][j], s[j] * P[j][i], abs_tol=tol):
                violations.append((i, j), s[i] * P[i][j], s[j] * P[j][i])
    return violations


viol = check_detailed_balance(P_rw, s_rw)
print("detailed balance violations:", len(viol))

P1 = np.array([
    [0.6, 0.4, 0.0],  # 状态0: 60%留在0, 40%转到1
    [0.0, 0.0, 1.0],  # 状态1: 100%转到2
    [0.5, 0.5, 0.0]  # 状态2: 50%转到0, 50%转到1
])


def simulate_chain(P, start, n_steps):
    """
    Simulate Markov chain with transition matrix P starting from state index start
    for n_steps.
    Returns the list of visted states including the start state.
    :param P:
    :param start:
    :param n_steps:
    :return:
    """
    states = [start]
    cur = start
    for _ in range(n_steps):
        cur = npr.choice(a=len(P), p=P[cur])
        states.append(cur)

    return states


print("\nVisialization:emperical distribution convergence for P1 and stationary distribution bar for random walk")
# plot empirical convergence for P1:running frequency over time for a single long run
traj = simulate_chain(P1, 0, 2000)
freqs = [Counter(traj[:t + 1]) for t in range(len(traj))]
frac0 = [freqs[t].get(0, 0) / (t + 1) for t in range(len(traj))]
frac1 = [freqs[t].get(1, 0) / (t + 1) for t in range(len(traj))]
frac2 = [freqs[t].get(2, 0) / (t + 1) for t in range(len(traj))]

plt.figure(figsize=(8, 4))
plt.plot(frac0, label='state 0')
plt.plot(frac1, label='state 1')
plt.plot(frac2, label='state 2')
plt.title("Empirical frequencies over time (P1 chain)")
plt.xlabel("time step")
plt.ylabel("empirical frequency")
plt.legend()
plt.tight_layout()
plt.show()

# plot stationary distribution for random walk
plt.figure(figsize=(6, 4))
plt.bar(range(len(s_rw)), s_rw)
plt.title("Stationary distribution for simple random walk (propotional to degrees)")
plt.xlabel("state (node)")
plt.ylabel("stationary prob")
plt.xticks(range(len(s_rw)))
plt.tight_layout()
plt.show()
